#include "lwipstack.h"
#include "socks_connect.h"

err_t driver2socks::LWIPStack::lwip_tcp_close(tcp_pcb *pcb) {
    if (pcb->callback_arg) {
        auto p = static_cast<SocksConnect*>(pcb->callback_arg);
        p->socks2tcp_nonfull.test_and_set();
        p->socks2tcp_nonfull.notify_all();
        pcb->callback_arg = nullptr;
        delete p;
    }

    LWIPStack::lwip_tcp_receive(pcb, NULL);
    err_t err = tcp_shutdown(pcb, 1, 1) | tcp_close(pcb);
    tcp_output(pcb);
    return err;
}

void driver2socks::LWIPStack::postTcpWrite(tcp_pcb *pcb, std::shared_ptr<NetPacket> buff) {
    if (pcb->state == 0 || TCP_STATE_IS_CLOSING(pcb->state) || pcb->callback_arg == nullptr) return;
    SocksConnect* ta = (SocksConnect*)(pcb->callback_arg);
    // ta->queue_socks2tcp.push(buff);
    if (!ta->queue_socks2tcp.try_push(buff)) {
        ta->socks2tcp_nonfull.clear();
        ta->socks2tcp_nonfull.wait(false);
    }

    StackEvent::Ptr e = std::make_shared<StackEvent>();
    e->pcb = pcb;
    e->type = TCP_WRITE;
    while(!queue_event_.try_push(e)) {
        queue_event_nonfull_.clear();
        queue_event_nonfull_.wait(false,std::memory_order_relaxed);
    }
}

bool driver2socks::LWIPStack::doWrite(tcp_pcb *pcb)
{
    if (pcb->state == CLOSED || TCP_STATE_IS_CLOSING(pcb->state) || pcb->callback_arg == nullptr) {
        return true;
    }
    SocksConnect* ta = (SocksConnect*)(pcb->callback_arg);
    err_t err = ERR_OK;

    while (!ta->queue_socks2tcp.empty()) {
        NetPacket::Ptr packet = *(ta->queue_socks2tcp.front());
        if (packet->data_len > tcp_sndbuf(pcb)) {
            err = tcp_output(pcb);
            if (err != ERR_OK) {
                SERROR("tcp_output fail:{}",err);
            }
            return false;
        }

        err = tcp_write(pcb, packet->data, packet->data_len, /*TCP_WRITE_FLAG_COPY*/0);
        if (err != ERR_OK) {
            SERROR("tcp_write error:{}",err);
            return false;
        }
        ta->queue_socks2tcp.pop();
        if (!ta->socks2tcp_nonfull.test()) {
            if (ta->queue_socks2tcp.size() < ta->queue_socks2tcp.capacity()*0.7) {
                ta->socks2tcp_nonfull.test_and_set();
                ta->socks2tcp_nonfull.notify_all();
            }
        }
        ta->queue_socks2tcp_wait_release.push(packet);
        //To prompt the system to send data now, call tcp_output() after calling tcp_write().
        err = tcp_output(pcb);
        if (err != ERR_OK) {
            SERROR("tcp_output fail:{}",err);
        }
    }
    return true;
}

err_t driver2socks::LWIPStack::onTcpRecv(void *arg, tcp_pcb *tpcb, pbuf *p, err_t err)
{
    if (tpcb == NULL) return ERR_VAL;

    if (p == NULL) {
        // SDEBUG("The remote host closed the connection.pcb({:#x})",(uint64_t)tpcb);
        LWIPStack::getInstance().lwip_tcp_close(tpcb);
        if (p) pbuf_free(p);
        return ERR_OK;
    }
    if (err != ERR_OK || 0 == p->len) {
        LWIPStack::getInstance().lwip_tcp_close(tpcb);
        return ERR_OK;
    }
    auto tcparg = static_cast<SocksConnect*>(tpcb->callback_arg);
    if (tcparg == nullptr) return ERR_OK;

    uint16_t data_len = p->tot_len;
    NetPacket::Ptr recv_data = _NetPacketPool->getSharedPacket(data_len);
    pbuf_copy_partial(p, recv_data->data, data_len, 0);
    pbuf_free(p);


    if (!tcparg->queue_tcp2socks.try_emplace(recv_data)) {
        SERROR("queue_tcp2socks try_emplace fail.");
        LWIPStack::getInstance().lwip_tcp_close(tpcb);
        return ERR_ABRT;
    }

    while (!tcparg->queue_tcp2socks.empty()) {
        NetPacket::Ptr np = *(tcparg->queue_tcp2socks.front());
        bool ret = tcparg->send2Socks(np);
        if (ret) {
            LWIPStack::getInstance().lwip_tcp_recved(tpcb, np->data_len);
            tcparg->queue_tcp2socks.pop();
            continue;
        } else {
            break;
        }
    }

    return ERR_OK;
}

err_t driver2socks::LWIPStack::onTcpAccept(void *arg, tcp_pcb *newpcb, err_t err)
{
    if (err != ERR_OK || newpcb == NULL) return ERR_VAL;
    auto ctx = new SocksConnect(newpcb);
    newpcb->callback_arg = ctx;
    ctx->setSocksServer(g_config->socks5_server_ip,g_config->socks5_server_port);
    ctx->start();

    return ERR_OK;
#if 0
    char strlocalip[64] = { 0 };
    uint32_t localip[4] = { 0 };
    if (newpcb->local_ip.type == IPADDR_TYPE_V4) {
        localip[0] = WinDivertHelperNtohl(newpcb->local_ip.u_addr.ip4.addr);
        WinDivertHelperFormatIPv4Address(localip[0], strlocalip, 64);
    } else {
        WinDivertHelperNtohIpv6Address(newpcb->local_ip.u_addr.ip6.addr, localip);
        WinDivertHelperFormatIPv6Address(localip, strlocalip, 64);
    }
    char strremoteip[64] = { 0 };
    uint32_t remoteip[4] = { 0 };
    if (newpcb->remote_ip.type == IPADDR_TYPE_V4) {
        remoteip[0] = WinDivertHelperNtohl(newpcb->remote_ip.u_addr.ip4.addr);
        WinDivertHelperFormatIPv4Address(remoteip[0], strremoteip, 64);
    } else {
        WinDivertHelperNtohIpv6Address(newpcb->remote_ip.u_addr.ip6.addr, remoteip);
        WinDivertHelperFormatIPv6Address(remoteip, strremoteip, 64);
    }
    SDEBUG("newtcp({:#x}) local ip:{},port:{},remote ip:{},port:{}",(uint64_t)newpcb,
           strlocalip,(int)newpcb->local_port,strremoteip,(int)newpcb->remote_port);


            //创建一个socks5代理客户端并关联tcp_pcb
            std::shared_ptr<driver2socks::SocksClient> context = driver2socks::SocksClient::make();
            TcpArg* targ = new TcpArg;
            targ->sc_client = context;
            newpcb->callback_arg = targ;
            context->lwip_pcb_ = newpcb;
            //设置tcp堆栈接收数据回调
            setTcpRecvCallback(newpcb);
            // LWIPStack::getInstance().setTcpSentCallback(newpcb);
            //tcp_nagle_disable(newpcb);

            auto callback_socksclient_recv = [](const asio::error_code& err, std::shared_ptr<NetPacket> buffer, std::size_t sz, tcp_pcb* tcp) {
                if (err) {
                    LWIPStack::getInstance().postTcpClose(tcp);
                    return;
                }

                LWIPStack::getInstance().postTcpWrite(tcp, buffer, (u16_t)sz, TCP_WRITE_FLAG_COPY,
                                                          [](err_t err) {
                                                              if (err != ERR_OK) {
                                                                  std::cerr << "tcp_write error:" << (int)err << "\n";
                                                              }
                                                          });

            };

            //建立代理通道
            char hoststr[64] = { 0 };
            uint32_t ip[4] = { 0 };
            if (newpcb->local_ip.type == IPADDR_TYPE_V4) {
                ip[0] = WinDivertHelperNtohl(newpcb->local_ip.u_addr.ip4.addr);
                WinDivertHelperFormatIPv4Address(ip[0], hoststr, 64);
            } else {
                WinDivertHelperNtohIpv6Address(newpcb->local_ip.u_addr.ip6.addr, ip);
                WinDivertHelperFormatIPv6Address(ip, hoststr, 64);
            }
            //开始socks5客户端（建立到服务器的连接并握手）
            context->start_socks(g_config->socks5_server_ip, g_config->socks5_server_port,
                                 std::string(hoststr), newpcb->local_port,
                                 callback_socksclient_recv,
                                 [](asio::error_code ec) {});
            return ERR_OK;
#endif
}

err_t driver2socks::LWIPStack::onTcpSent(void *arg, tcp_pcb *tpcb, u16_t len)
{
    // if (arg == nullptr) return ERR_OK;

    SocksConnect* sc = (SocksConnect*)(tpcb->callback_arg);
    if (sc == nullptr) return ERR_OK;

    auto queue = &(sc->queue_socks2tcp_wait_release);

    sc->remaining += len;
    while (sc->remaining > 0 && !queue->empty()) {
        NetPacket::Ptr net_pack = *(queue->front());
        if (sc->remaining >= net_pack->data_len) {
            sc->remaining -= net_pack->data_len;
            queue->pop();
        } else {
            break;
        }
    }
    return ERR_OK;
}
